{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Object Detection Model Training using Synthetic Data from Omniverse Replicator Objects Extension\n",
    "\n",
    "[NVIDIA Omniverse™ Isaac Sim](https://docs.omniverse.nvidia.com/isaacsim/latest/index.html) is a robotics simulation toolkit for the NVIDIA Omniverse™ platform. Isaac Sim has essential features for building virtual robotic worlds and experiments. It provides researchers and practitioners with the tools and workflows they need to create robust, physically accurate simulations and synthetic datasets.\n",
    "\n",
    "## What is Replicator Object Extension?\n",
    "\n",
    "[omni.replicator.object](https://docs.omniverse.nvidia.com/isaacsim/latest/replicator_tutorials/tutorial_replicator_object.html) is an extension that generates synthetic data for model training without requiring any changes to the code. It can be used for various tasks, such as retail object detection and robotics. The extension takes a YAML description file as input, which describes a mutable scene or a hierarchy of stacked description files. It then outputs a description file along with graphics content, including RGB images, 2D/3D bounding boxes, segmentation masks, and more.\n",
    "\n",
    "<img align=\"center\" src=\"https://docs.omniverse.nvidia.com/isaacsim/latest/_images/overview.png\" width=\"540\">\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to generate the synthetic dataset for DINO TAO training:\n",
    "\n",
    "* Setup a configuration file for generating synthetic dataset in Omniverse.\n",
    "* Generate the synthetic training data based on the configuration file.\n",
    "* Visualize the generated images along with their annotations.\n",
    "* Train an object detection model using freshly generated synthetic data with TAO\n",
    "\n",
    "At the end of this notebook, you will have generated a synthetic training set and an object detection model which you may deploy via [TAO-Deploy](https://github.com/NVIDIA/tao_deploy)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the generated synthetic dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/synthetic_dino/results`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\", \"synthetic_dino\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"synthetic_dino\", \"results\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/synthetic_dino\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tao_configs = {\n",
    "   \"Mounts\":[\n",
    "         # Mapping the Local project directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         },\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid()),\n",
    "        \"network\": \"host\"\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tao_configs, mfile, indent=4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in the `nvidia-pyindex` python index. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-pyindex\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Generate synthetic dataset and verify the annotations <a class=\"anchor\" id=\"head-2\"></a>\n",
    "### 2.1 Setup the configuration file for the Omniverse Replicator\n",
    "In this section, we will use the \"pallet\" dataset for the tutorial. The following script will automatically generate the synthetic \"pallet\" dataset:\n",
    "\n",
    "- Please note that all the USD contents are sourced from `omniverse://content.ov.nvidia.com/`.\n",
    "- Make sure you have the necessary access before running the code.\n",
    "- More details for the hyper-parameters can be found in the [omni.replicator.object](https://docs.omniverse.nvidia.com/isaacsim/latest/replicator_tutorials/tutorial_replicator_object.html) page.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file=\"\"\"\n",
    "omni.replicator.object:\n",
    "  version: 0.2.16\n",
    "  num_frames: 20\n",
    "  seed: 100\n",
    "  inter_frame_time: 1\n",
    "  gravity: 10000\n",
    "  position_H:\n",
    "    harmonizer_type: mutable_attribute\n",
    "    mutable_attribute:\n",
    "      distribution_type: range\n",
    "      start:\n",
    "      - -94.77713317047056\n",
    "      - 0\n",
    "      - -35.661244451558446\n",
    "      end:\n",
    "      - -94.77713317047056\n",
    "      - 0\n",
    "      - -35.661244451558446\n",
    "  screen_height: 720\n",
    "  focal_length: 14.228393962367306\n",
    "  output_path: /tmpsrc/results\n",
    "  horizontal_aperture: 20.955\n",
    "  screen_width: 1080\n",
    "  camera_parameters:\n",
    "    far_clip: 100000\n",
    "    focal_length: $(focal_length)\n",
    "    horizontal_aperture: $(horizontal_aperture)\n",
    "    near_clip: 0.1\n",
    "    screen_height: $(screen_height)\n",
    "    screen_width: $(screen_width)\n",
    "  default_camera:\n",
    "    count: 1\n",
    "    camera_parameters: $(camera_parameters)\n",
    "    transform_operators:\n",
    "    - translate_global:\n",
    "        distribution_type: harmonized\n",
    "        harmonizer_name: position_H\n",
    "    - rotateY: $[seed]*20\n",
    "    - rotateX:\n",
    "        distribution_type: range\n",
    "        start: -15\n",
    "        end: -25\n",
    "    - translate:\n",
    "        distribution_type: range\n",
    "        start:\n",
    "        - -40\n",
    "        - -30\n",
    "        - 400\n",
    "        end:\n",
    "        - 40\n",
    "        - 30\n",
    "        - 550\n",
    "    type: camera\n",
    "  distant_light:\n",
    "    color:\n",
    "      distribution_type: range\n",
    "      end:\n",
    "      - 1.3\n",
    "      - 1.3\n",
    "      - 1.3\n",
    "      start:\n",
    "      - 0.7\n",
    "      - 0.7\n",
    "      - 0.7\n",
    "    count: 5\n",
    "    intensity:\n",
    "      distribution_type: range\n",
    "      end: 600\n",
    "      start: 150\n",
    "    subtype: distant\n",
    "    transform_operators:\n",
    "    - rotateY:\n",
    "        distribution_type: range\n",
    "        end: 180\n",
    "        start: -180\n",
    "    - rotateX:\n",
    "        distribution_type: range\n",
    "        end: -10\n",
    "        start: -40\n",
    "    type: light\n",
    "  dome_light:\n",
    "    type: light\n",
    "    subtype: dome\n",
    "    color:\n",
    "      distribution_type: range\n",
    "      start:\n",
    "      - 0.7\n",
    "      - 0.7\n",
    "      - 0.7\n",
    "      end:\n",
    "      - 1.3\n",
    "      - 1.3\n",
    "      - 1.3\n",
    "    intensity:\n",
    "      distribution_type: range\n",
    "      start: 1000\n",
    "      end: 3000\n",
    "    transform_operators:\n",
    "    - rotateX: 270\n",
    "  plane:\n",
    "    physics: collision\n",
    "    type: geometry\n",
    "    subtype: plane\n",
    "    tracked: false\n",
    "    transform_operators:\n",
    "    - scale:\n",
    "      - 5\n",
    "      - 5\n",
    "      - 5\n",
    "  rotY_H:\n",
    "    harmonizer_type: mutable_attribute\n",
    "    mutable_attribute:\n",
    "      distribution_type: range\n",
    "      start: 0\n",
    "      end: 0\n",
    "  translate_H:\n",
    "    harmonizer_type: mutable_attribute\n",
    "    mutable_attribute:\n",
    "      distribution_type: range\n",
    "      start:\n",
    "      - 0\n",
    "      - 60\n",
    "      - 0\n",
    "      end:\n",
    "      - 0\n",
    "      - 30\n",
    "      - 0\n",
    "\n",
    "  box:\n",
    "    count: 3\n",
    "    physics: rigidbody\n",
    "    type: geometry\n",
    "    subtype: mesh\n",
    "    tracked: true\n",
    "    transform_operators:\n",
    "    - translate_global:\n",
    "        distribution_type: harmonized\n",
    "        harmonizer_name: position_H\n",
    "    - translate_pallet:\n",
    "        distribution_type: harmonized\n",
    "        harmonizer_name: translate_H\n",
    "    - rotateY:\n",
    "        distribution_type: harmonized\n",
    "        harmonizer_name: rotY_H\n",
    "    - translate:\n",
    "      - 120 * ($[index])\n",
    "      - 20\n",
    "      - 0\n",
    "    - rotateXYZ:\n",
    "      - 0\n",
    "      - -90\n",
    "      - -90\n",
    "    - scale:\n",
    "      - 12\n",
    "      - 10\n",
    "      - 6\n",
    "    usd_path:\n",
    "      distribution_type: set\n",
    "      values:\n",
    "      - omniverse://content.ov.nvidia.com/NVIDIA/Assets/DigitalTwin/Assets/Warehouse/Shipping/Cardboard_Boxes/White_A/WhiteCorrugatedBox_A01_10x10x10cm_PR_NVD_01.usd\n",
    "      - omniverse://content.ov.nvidia.com/NVIDIA/Assets/DigitalTwin/Assets/Warehouse/Shipping/Cardboard_Boxes/Cube_A/CubeBox_A01_10cm_PR_NVD_01.usd\n",
    "  warehouse:\n",
    "    type: geometry\n",
    "    subtype: mesh\n",
    "    usd_path: omniverse://content.ov.nvidia.com/NVIDIA/Assets/Isaac/2023.1.1/Isaac/Environments/Simple_Warehouse/warehouse_with_forklifts.usd\n",
    "    transform_operators:\n",
    "    - translate:\n",
    "      - -200\n",
    "      - 0.1\n",
    "      - 0\n",
    "    - rotateXYZ:\n",
    "      - 0\n",
    "      - -90\n",
    "      - -90\n",
    "    - scale:\n",
    "      - 100\n",
    "      - 100\n",
    "      - 100\n",
    "\n",
    "  output_switches:\n",
    "    images: True\n",
    "    labels: True\n",
    "    descriptions: False\n",
    "    3d_labels: False\n",
    "    segmentation: False\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the configuration file and save to the local.\n",
    "import yaml\n",
    "yaml_file = yaml.safe_load(config_file)\n",
    "with open(os.path.join(os.getenv(\"HOST_DATA_DIR\", os.getcwd()), 'config.yaml'), 'w') as outfile:\n",
    "    yaml.dump(yaml_file, outfile, default_flow_style=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Note that you could simple adjust the hyper-parameter from config file to generate the synthetic data on your use case.\n",
    "* **If you are using content from omniver nucleus servers you need to set your own OMNI_USER and OMNI_PASS variables.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the Omniverse password\n",
    "%env OMNI_PASS=YOUR_OWN_OMNI_PASSWORD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Launch the synthetic data generation\n",
    "Launch the synthetic data generation with the Omniverse replicator object extension inside the container.\n",
    "\n",
    "The following example is one of the synthetic data generation scenes. You can modify the config file to generate different scenes with various objects, backgrounds, and numbers of target objects.\n",
    "* Note that the current synthetic data generation pipeline only supports single GPU. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker run --gpus device=0 -it \\\n",
    "    --entrypoint /bin/bash \\\n",
    "    --network host \\\n",
    "    -v $HOST_DATA_DIR:/tmpsrc \\\n",
    "    -e OMNI_USER='$omni-api-token' \\\n",
    "    -e OMNI_PASS=$OMNI_PASS \\\n",
    "    nvcr.io/nvidia/isaac-sim:4.0.0 \\\n",
    "    -c \"apt-get update && apt-get install libglib2.0-dev -y && bash isaac-sim.sh --no-window --allow-root --/windowless=True --allow-root --/log/outputStreamLevel=fatal --/app/extensions/fastImporter/enabled=false --enable omni.replicator.object --/config/file=/tmpsrc/config.yaml --/binary=True\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Visualize the generated data\n",
    "In this section, we run the `visualization` tool to visualize the generated synthetic data. The generation tool produces the synthetic images and the corresponding json files that contain the training annotation information.\n",
    "\n",
    "After generating the synthetic data, it stores in the \"HOST_RESULTS_DIR\". It can be visualized by following commands. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the matplotlib dependencies.\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\"\n",
    "!pip3 install opencv-python==4.8.0.74\n",
    "!pip3 install numpy==1.24.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg']\n",
    "\n",
    "def visualize_images(output_path, num_cols=4, num_images=10):\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[40,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "# Note that the sample spec is not meant to produce SOTA (state-of-the-art) accuracy on Objectron dataset.\n",
    "IMAGE_DIR = os.path.join(os.environ['HOST_DATA_DIR'], 'results', 'images')\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(IMAGE_DIR, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Visualize and Verify the Training Data Annotation \n",
    "In this section, we visualize the annotation for the generated synthetic data by using the related json file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "import os\n",
    "import cv2\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg']\n",
    "\n",
    "def visualize_images_with_annotations(images_path, labels_path, num_cols=4, num_images=10):\n",
    "    \n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[40,30])\n",
    "    f.tight_layout()\n",
    "    \n",
    "    images=[]\n",
    "    labels=[]\n",
    "    for label_file in os.listdir(labels_path):\n",
    "        labels.append(os.path.join(labels_path, label_file))\n",
    "        images.append(os.path.join(images_path, label_file.replace('.txt', '.jpg')))\n",
    "        \n",
    "    for idx, img_path in enumerate(images[:num_images]):\n",
    "        img = plt.imread(img_path)\n",
    "        lf = labels[idx]\n",
    "        with open(labels[idx], 'r') as fp:\n",
    "            for line in fp:\n",
    "                sp = line.split(' ')\n",
    "                xmin, ymin, xmax, ymax = int(sp[4]), int(sp[5]), int(sp[6]), int(sp[7])\n",
    "                cv2.rectangle(img, (xmin,ymin), (xmax,ymax), (0,255,0), 2)\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images along with their annotations\n",
    "IMAGE_DIR = os.path.join(os.environ['HOST_DATA_DIR'], 'results', 'images')\n",
    "LABELS_DIR = os.path.join(os.environ['HOST_DATA_DIR'], 'results', 'labels')\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images_with_annotations(IMAGE_DIR, LABELS_DIR, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Launch the DINO TAO object detection training using the generated synthetic dataset\n",
    "In this section, it will introduce how to use the synthetic dataset to launch the training. \n",
    "More details regarding to the hyper-parameters and the end-to-end training pipeline could be found in the [DINO Notebook](https://github.com/NVIDIA/tao_tutorials/blob/main/notebooks/tao_launcher_starter_kit/dino/dino.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Download the pre-trained model\n",
    "We will use NGC CLI to get the pre-trained models. For more details, go to [ngc.nvidia.com](ngc.nvidia.com) and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!mkdir -p $HOST_RESULTS_DIR/pretrained_models\n",
    "!ngc registry model download-version nvidia/tao/retail_object_detection:trainable_binary_v2.1.1 --dest $HOST_RESULTS_DIR/pretrained_models\n",
    "print(\"Check if model is downloaded into dir.\")\n",
    "!ls -l $HOST_RESULTS_DIR/pretrained_models/retail_object_detection_vtrainable_binary_v2.1.1/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    "\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * train_data_sources:\n",
    "        * image_dir: the root directory for train images\n",
    "        * json_file: annotation file for train data. Must be in COCO JSON format.\n",
    "    * val_data_sources: \n",
    "        * image_dir: the root directory for validation images\n",
    "        * json_file: annotation file for validation data. Must be in COCO JSON format.\n",
    "    * num_classes: number of classes of your training data\n",
    "    * batch_size: batch size for dataloader\n",
    "    * workers: number of workers to do data loading\n",
    "* model: configure the model setting\n",
    "    * pretrained_backbone_path: path to the pretrained backbone model. ResNet50, FAN-variants, and GCViT-variants are supported\n",
    "    * num_feature_levels: number of feature levels used from backbone\n",
    "    * dec_layers: number of decoder layers\n",
    "    * enc_layers: number of encoder layers\n",
    "    * num_queries: number of queries for the model\n",
    "    * num_select: number of top-k proposals to select from\n",
    "    * use_dn: flag to enable denoising during training\n",
    "    * dropout_ratio: drop out ratio\n",
    "* train: configure the training hyperparameters\n",
    "    * pretrained_model_path: load pretrained model path before train\n",
    "    * freeze: freezes listed modules dutraining train\n",
    "    * num_gpus: number of gpus \n",
    "    * num_nodes: number of nodes (num_nodes=1 for single node)\n",
    "    * val_interval: validation interval\n",
    "    * optim:\n",
    "        * lr_backbone: learning rate for backbone\n",
    "        * lr: learning rate for the rest of the model\n",
    "        * lr_steps: learning rate decay step milestone (MultiStep)\n",
    "    * num_epochs: number of epochs\n",
    "    * activation_checkpoint: recompute activations in the backward to save GPU memory. Default is `True`.\n",
    "    * precision: If set to fp16, the training is run on Automatic Mixed Precision (AMP)\n",
    "    * distributed_strategy: Default is `ddp`. `ddp_sharded` is also supported.\n",
    "\n",
    "See the [TAO documentation - DINO](https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/dino.html) to get all the parameters that are configurable.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/train_synthetic.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Run object detection TAO training <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* Evaluation uses COCO metrics. For more info, please refer to: https://cocodataset.org/#detection-eval\n",
    "\n",
    "* To speed up training, try setting `train.precision=fp16` for mixed precision training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results\n",
    "%env MODEL_DIR = /model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"For multi-GPU, change num_gpus in train.yaml based on your machine.\")\n",
    "print(\"For multi-node, change num_gpus and num_nodes in train.yaml based on your machine.\")\n",
    "# If you face out of memory issue, you may reduce the batch size in the spec file by passing dataset.batch_size=2\n",
    "!tao model dino train \\\n",
    "          -e $SPECS_DIR/train_synthetic.yaml \\\n",
    "          results_dir=$RESULTS_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Trained checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"train/dino_model_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/train/retail_object_detection.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/retail_object_detection.pth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has come to an end.\n",
    "More details of end-to-end training and inference pipeline could be found in the [DINO Notebook](hhttps://github.com/NVIDIA/tao_tutorials/blob/main/notebooks/tao_launcher_starter_kit/dino/dino.ipynb) and [Retail Object Detection Notebook](https://github.com/NVIDIA/tao_tutorials/blob/main/notebooks/tao_launcher_starter_kit/retail_object_detection/retail_object_detection.ipynb).\n",
    "\n",
    "Note that the model we trained above is not a competitive object detection model. We have just trained few epochs with very few data without any pre-trained weights. Our goal was to show how to generate synthetic data using the extension and train an object detection model using those images. The training shows that the training error is decreasing so we basically achieved our goal."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env_notebook5",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
